import torch
# from pytorch_msssim import ms_ssim
import os
import random
from tqdm import tqdm
from torch.utils.data import DataLoader
from dataset import MyDataset
from vae import VAE, Discriminator
from loss import VaeLossFn, WGANGP
from pred import prediction


if __name__ == '__main__':
    img_rootPath = r'CASIA-WebFace'
    saveModel_dir = r'checkpoints'
    verification = r'verification'
    perceptualLoss_weight = r'vgg16-397923af.pth'
    load_model = r''
    epoch = 1000
    batch_size = 128
    d_lr = 1e-5
    g_lr = 1e-3
    
        
    train_ds = MyDataset(img_rootPath, 'train', 0.98)
    test_ds = MyDataset(img_rootPath, 'test', 0.98)
    train_dl = DataLoader(train_ds, shuffle=True, batch_size=batch_size, drop_last=True)
    test_dl = DataLoader(test_ds, shuffle=True, batch_size=batch_size, drop_last=True)


    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = VAE()
    if load_model:
        model.load_state_dict(torch.load(load_model), strict=False)
    model.to(device)
    discriminator = Discriminator()
    discriminator.to(device)

    optim = torch.optim.Adam(model.parameters(), g_lr, betas=(0.5, 0.999))
    gan_optim = torch.optim.Adam(discriminator.parameters(), d_lr, betas=(0.5, 0.999))
    vae_loss_fn = VaeLossFn(perceptualLoss_weight, device)
    gan_loss_fn = WGANGP(device)
    maxAcc_test = 0
    
    for e in range(epoch):
        model.train()
        with tqdm(train_dl) as train_dl_buffer:
            
            gan_total_loss = 0
            main_total_loss = 0
            vae_loss_fn.init_kmean_each_epoch()
            for batch, img in enumerate(train_dl_buffer):
                img = img.to(device)

                # gan对抗器训练 ---------------------------------------------------------------------------------#
                with torch.no_grad():
                    out_img, P = model(img)

                real = discriminator(img)
                fake = discriminator(out_img)

                loss = gan_loss_fn.discriminator_loss(real, fake, img, out_img, discriminator)
                
                gan_optim.zero_grad()
                loss.backward()
                gan_optim.step()
                gan_total_loss += loss.item()

                # gan生成器训练 ---------------------------------------------------------------------------------#
                out_img, P = model(img)

                fake = discriminator(out_img)
                vae_loss = vae_loss_fn(out_img, img, P)
                gan_loss = gan_loss_fn.generator_loss(fake)
                
                loss = 5e-5*gan_loss + vae_loss
                # loss = vae_loss

                optim.zero_grad()
                gan_optim.zero_grad()
                loss.backward()
                optim.step()
                main_total_loss += loss.item()
                
                train_dl_buffer.set_description("Epoch {}: vaeloss {:.8f} discriminatorloss {:.8f}".format(e+1, main_total_loss/(batch+1), gan_total_loss/(batch+1)))
                # train_dl_buffer.set_description("Epoch {}: vaeloss {:.8f}".format(e+1, main_total_loss/(batch+1)))

        # 保存模型 
        path = os.path.join(saveModel_dir, 'checkpoint.pth')
        torch.save(model.state_dict(), path)
        print("保存模型===》{}".format(path))

        # # 模型测试
        # total_acc = 0
        # model.eval()
        # with tqdm(test_dl) as test_dl_buffer:
        #     for batch, (img, mask, target) in enumerate(test_dl_buffer):

        #         with torch.no_grad():
        #             out = model(img.to(device), mask.to(device))
        #             acc = ms_ssim(out, target.to(device), data_range=1, size_average=True)
                
        #         total_acc += acc.item()
        #         batch += 1
        #         test_dl_buffer.set_description("Epoch {}: testAcc {:.4f}".format(e+1, total_acc/batch))
        
        # if total_acc/batch > maxAcc_test:
        #     maxAcc_test = total_acc/batch
        #     path = os.path.join(saveModel_dir, 'best_test.pth')
        #     torch.save(model.state_dict(), path)
        #     print("保存模型===》{}".format(path))
        
        for i in range(10):
            img_path = test_ds.img_path_arr[random.randint(0, test_ds.__len__()-1)]
            save_path = os.path.join(verification, f'{i}.jpg')
            prediction(model, img_path, save_path, device)